import torch
import torch.nn as nn
import torch.nn.functional as F


class STGCNLayer(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(STGCNLayer, self).__init__()
        self.temporal1 = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 3), padding=(0, 1))
        self.graph_conv = nn.Linear(out_channels, out_channels)
        self.temporal2 = nn.Conv2d(out_channels, out_channels, kernel_size=(1, 3), padding=(0, 1))

    def forward(self, x, adj):
        """
        x: (batch_size, in_channels, num_nodes, seq_len)
        adj: (num_nodes, num_nodes)
        """
        # Temporal Convolution 1
        x = self.temporal1(x)  # shape: (batch, out_channels, num_nodes, seq_len)

        # Graph Convolution
        batch_size, out_channels, num_nodes, seq_len = x.shape
        x = x.permute(0, 3, 2, 1)  # -> (batch, seq_len, num_nodes, out_channels)
        x = torch.matmul(adj, x)  # 图卷积传播 -> (batch, seq_len, num_nodes, out_channels)
        x = self.graph_conv(x)  # -> Linear 层 -> (batch, seq_len, num_nodes, out_channels)
        x = x.permute(0, 3, 2, 1)  # -> (batch, out_channels, num_nodes, seq_len)

        # Temporal Convolution 2
        x = self.temporal2(x)
        return F.relu(x)


class STGNN(nn.Module):
    def __init__(self, num_nodes, in_channels, hidden_channels, out_channels):
        super(STGNN, self).__init__()
        self.gconv1 = STGCNLayer(in_channels, hidden_channels)
        self.gconv2 = STGCNLayer(hidden_channels, out_channels)
        self.pool = nn.AdaptiveAvgPool2d((num_nodes, 1))  # 池化为每个节点一个输出
        self.classifier = nn.Linear(out_channels, 1)  # 二分类 or 回归

    def forward(self, x, adj):
        x = self.gconv1(x, adj)
        x = self.gconv2(x, adj)
        x = self.pool(x)  # shape: (batch_size, out_channels, num_nodes, 1)
        x = x.squeeze(-1).permute(0, 2, 1)  # -> (batch_size, num_nodes, out_channels)
        out = self.classifier(x)  # -> (batch_size, num_nodes, 1)
        return out.squeeze(-1)  # -> (batch_size, num_nodes)
